package com.gmall.utils

import com.gmall.bean.AlertInfo
import io.searchbox.client.{JestClient, JestClientFactory}
import io.searchbox.client.config.HttpClientConfig
import io.searchbox.core.{Bulk, Index}
import org.apache.spark.rdd.RDD

import scala.collection.JavaConverters._

/**
 * <p>文件名称: ES工具类 </p>
 * <p>描述: 单个文档插入、批量文档插入 </p>
 * <p>创建时间: 2020/10/21</p>
 * @author <a href="mail to: 1034652894@qq.com" rel="nofollow">Demik Qin</a>
 * @version v1.0
 * @update [序号][日期YYYY-MM-DD] [更改人姓名][变更描述]
 */
object MyESUtil {
    // 获取ES的集群URL地址，需要转成java集合
    private val ESUrl: String = PropertiesUtil.getValue("ESUrl")
    // 创建ES客户端工厂
    private val factory: JestClientFactory = new JestClientFactory
    // 构建客户端配置文件
    private val config: HttpClientConfig = new HttpClientConfig.Builder(ESUrl.split(",").toList.asJava) //配置URL地址
      .maxTotalConnection(100) //最大连接数
      .connTimeout(10000) //连接超时时间
      .readTimeout(10000) //读取数据超时时间
      .build()
    // 加载客户端配置文件
    factory.setHttpClientConfig(config)
    case class User(name: String, age: Int)
    def main(args: Array[String]): Unit = {
        val json =
            """
              |
              |{
              | "user":"taylor",
              | "age":20
              |}
              |""".stripMargin
              new User("a",20)
        val users = Iterator(User("a", 20), User("b", 21), User("c", 22), User("d", 23))
        insertBulk("text", users)
    }

    /**
     * 向ES插入单条数据
     * @param indexName index的名称
     * @param source 具体的文档数据
     * @param id 文档的唯一编号，不传默认null，由ES自动生成
     */
    def insertSingle(indexName: String, source: Object, id: String = null): Unit = {
        val client: JestClient = factory.getObject
        // 创建一个document
        val action: Index = new Index.Builder(source)
          .index(indexName)
          .`type`("_doc")
          .id(id)
          .build()
        // 执行该文档
        client.execute(action)
        client.close()
    }


    /**
     * 向ES批次插入数据 sources: Iterable[(id, caseClass)] 或者 Iterable[caseClass]
     * @param indexName 批量插入的index名称
     * @param sources 批量插入的文档数据，是Iterator类型
     */
    def insertBulk(indexName: String, sources: Iterator[Object]): Unit = {
        val client: JestClient = factory.getObject
        val bulkBuilder = new Bulk.Builder()
          .defaultIndex(indexName) //批量插入的index名称
          .defaultType("_doc")
        // 遍历sources，把每个source都转换成一个文档，并添加到bulk中
        sources.foreach {
            // 如果传入的是二元元组, 第一个值表示id
            case (id: Any, source) =>
                bulkBuilder.addAction(new Index.Builder(source).id(id.toString).build())
            // 其他类型 没有id, ES会自动生成默认id
            case source =>
                bulkBuilder.addAction(new Index.Builder(source).build())
        }
        client.execute(bulkBuilder.build())
        client.close()
    }

    /**
     * 将AlertInfo类型RDD隐式转换，添加saveToES方法，
     * @param rdd AlertInfo类型
     */
    implicit class RichES(rdd: RDD[AlertInfo]) {
        def saveToES(indexName: String): Unit = {
            rdd.foreachPartition((iter: Iterator[AlertInfo]) => {
                MyESUtil.insertBulk(
                    indexName,
                    // 对同一个mid, 每分钟最多预警一次.
                    // 所以需要做去重处理.
                    // 去重使用es来实现去重: 当id相同的时候, 后面的会自动覆盖前面的
                    iter.map(info => (info.mid + "_" + info.ts / 1000 / 60, info)))
            })
        }
    }

}
